package LibDL.recommender.utils;

import LibDL.core.*;
import LibDL.core.nn.Embedding;

import java.math.BigInteger;
import java.util.*;

//activation_getter = {'iden': lambda x: x, 'relu': F.relu, 'tanh': torch.tanh, 'sigm': torch.sigmoid}

public class RecommenderUtils {

    public static Random random = new Random();

//    public Tensor gpu(Tensor tensor, boolean gpu) {
//        if (gpu) {
//            return tensor.cuda();
//        } else
//            return tensor;
//    }
//
//    public Tensor cpu(Tensor tensor) {
//        if (tensor.is_cuda) {
//            return tensor.cpu();
//        } else
//            return tensor;
//    }

    public static ArrayList<ArrayList<Tensor>> minibatch(Tensor[] tensors, int batch_size) {

        int len = tensors[0].size(0);

        ArrayList<ArrayList<Tensor>> arrayListArrayList = new ArrayList<>();
        for (int i = 0; i < len; i += batch_size) {
            ArrayList<Tensor> tensorArrayList = new ArrayList<>();
            for (Tensor tensor : tensors) {
                int nlen = batch_size;
                int t = i + batch_size - len;
                if (t > 0)
                    nlen -= t;
                tensorArrayList.add(tensor.narrow_copy(0, i, nlen));
            }
            arrayListArrayList.add(tensorArrayList);
        }

        return arrayListArrayList;
    }

    static Tensor[] shuffle(Tensor[] arrays) throws Exception {
        //boolean require_indices = false; // kwarg

        int t = arrays[0].size(0);
        for (Tensor arr : arrays) {
            if (arr.size(0) != t) {
                throw new Exception("All inputs to shuffle must have the same length.");
            }
        }

        Tensor shuffle_indices = functional.randperm(t, new TensorOptions((Dtype.INT64)));

        Tensor[] result = new Tensor[arrays.length];

        for (int i = 0; i < arrays.length; i++) {
            result[i] = arrays[i].index_select(0, shuffle_indices);
        }

        return result;
    }

    public void assert_no_grad(Tensor variable) {
        if (variable.requires_grad()) {
            System.out.println("nn criterions don't compute the gradient w.r.t. targets - please ");
            System.out.println("mark these variables as volatile or not requiring gradients");
        }
    }

    public static void set_seed(long seed, boolean cuda) {

        random.setSeed(seed);
//        if (cuda) //TODO CUDA
//            functional.cuda.manual_seed(seed);
//        else
        functional.manual_seed(BigInteger.valueOf(seed));
    }
}
